#!/usr/bin/env python
# -*- coding: utf-8 -*-


import zmq
import argparse
import logging

class ProxyClient:
    def __init__(self, context, host, req_port, sub_port, timeout=5000):
        self.host = host
        self.req_port = req_port
        self.sub_port = sub_port
        self.context = context
        self.timeout = timeout
        self.device_list = []
        # 创建REQ socket用于请求-响应模式
        self.req_socket = self.context.socket(zmq.REQ)
        self.req_socket.setsockopt(zmq.RCVTIMEO, timeout)
        self.req_socket.setsockopt(zmq.SNDTIMEO, timeout)
        self.req_address = f"tcp://{host}:{req_port}"
        # 创建SUB socket用于接收设备更新通知
        self.sub_socket = self.context.socket(zmq.SUB)
        self.sub_socket.setsockopt_string(zmq.SUBSCRIBE, "")
        self.sub_address = f"tcp://{host}:{sub_port}"
        # 连接到代理服务器
        self.req_socket.connect(self.req_address)
        logging.debug(f"REQ socket connected to {self.req_address}")
        self.sub_socket.connect(self.sub_address)
        logging.debug(f"SUB socket connected to {self.sub_address}")
    
    def __del__(self):
        self.close()

    def close(self):
        if self.req_socket:
            self.req_socket.close()
            self.req_socket = None
        if self.sub_socket:
            self.sub_socket.close()
            self.sub_socket = None

    def send_json(self, request_data):
        try:
            self.req_socket.send_json(request_data)
            return None
        except zmq.error.ZMQError as e:
            logging.error(f"发送JSON数据时发生ZMQ错误: {e}")
            return str(e)
        except TypeError as e:
            logging.error(f"发送JSON数据时发生类型错误: {e}")
            return str(e)
        except Exception as e:
            logging.error(f"发送JSON数据时发生未知错误: {e}")
            return str(e)

    def recv_json(self):
        try:
            response = self.req_socket.recv_json()
            return response
        except zmq.error.ZMQError as e:
            logging.error(f"发送JSON数据时发生ZMQ错误: {e}")
            return {"retval": 1, "message": f"接收JSON数据时发生ZMQ错误: {str(e)}"}
                
    def update_device_list(self) -> bool:
        # 发送获取设备列表请求
        request_data = {"api": "get_device_list"}
        errmsg = self.send_json(request_data)
        if errmsg is not None:
            logging.error(f"client {self.host} 获取设备列表失败: {errmsg}")
            return False
        response = self.recv_json()
        if response["retval"] == 0:
            self.device_list = response["result"]["devices"]
            for device in self.device_list:
                device["name"] = device["sn"]
                device["proxy"] = self.host
            logging.info(f"client {self.host} 成功获取设备列表: {self.device_list}")
            return True
        else:
            errmsg = response.get('message', '未知错误')
            logging.error(f"获取设备列表失败: {errmsg}")
        return False

class ProxyClientMgr:
    def __init__(self, proxy_configs, default_timeout=5000):
        self.context = zmq.Context()
        self.clients = {}
        self.device_clients = {}
        self.default_timeout = default_timeout
        for config in proxy_configs:
            name = config.get("NAME")
            host = config.get("HOST")
            req_port = config.get("REQ_PORT")
            sub_port = config.get("SUB_PORT")
            client = ProxyClient(self.context, host, req_port, sub_port, timeout=self.default_timeout)
            if client.update_device_list():
                self.clients[host] = client
                for device in client.device_list:
                    self.device_clients[device['sn']] = client
                logging.info(f"Successfully initialized ProxyClient: at {host}")
            else:
                logging.error(f"Failed to initialize ProxyClient: at {host}")
                client.close()
                del client

    def __del__(self):
        self.close_all()
        self.context.term()
        self.context = None
    
    def close_all(self):
        """关闭所有代理客户端连接"""
        for client in self.clients.values():
            try:
                client.close()
                logging.info(f"成功关闭代理客户端连接: {client.host}")
            except Exception as e:
                logging.error(f"关闭代理客户端连接时发生错误: {str(e)}")
        self.clients.clear()
        self.device_clients.clear()

    def send_request(self, reqquest_data):
        if reqquest_data['api'] == 'get_device_list':
            device_list = []
            for client in self.clients.values():
                device_list.extend(client.device_list)
            return {"retval": 0, "result": {"devices": device_list}}
        device_sn = reqquest_data['device_sn']
        client = self.device_clients.get(device_sn)
        if client is None:
            return {"retval": 1, "message": f"设备 {device_sn} 的代理不存在"}
        errmsg = client.send_json(reqquest_data)
        if errmsg is not None:
            logging.error(f"client {client.host} 发送请求失败: {errmsg}")
            return {"retval": 1, "message": f"client {client.host} 发送请求失败: {errmsg}"}
        return client.recv_json()

    def get_device(self, device_sn):
        client = self.device_clients.get(device_sn)
        if client is None:
            logging.error(f"设备 {device_sn} 的代理不存在")
            return None
        for device in client.device_list:
            if device['sn'] == device_sn:
                return device
        logging.error(f"设备 {device_sn} 不存在")
        return None

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="ProxyClient Example")
    parser.add_argument("-n", "--name", type=str, default="client", help="Client name")
    parser.add_argument("-H", "--host", type=str, default="172.16.1.84", help="Host address")
    parser.add_argument("--req-port", type=int, default=5555, help="REQ port")
    parser.add_argument("--sub-port", type=int, default=5556, help="SUB port")
    parser.add_argument("--timeout", type=int, default=5000, help="Timeout in milliseconds")
    args = parser.parse_args()
    logging.basicConfig(level=logging.DEBUG)
    zmq_context = zmq.Context()
    client = ProxyClient(zmq_context, args.host, args.req_port, args.sub_port, args.timeout)
    logging.info(f"device list: {client.get_device_list(force_refresh=True)}")
    client.close()
    zmq_context.term()